import numpy as np
from sklearn.metrics import precision_score, recall_score, f1_score


def calculate_metrics(y_true, y_pred):
    # 计算评估指标
    precision = precision_score(y_true, y_pred, average='weighted')
    recall = recall_score(y_true, y_pred, average='weighted')
    f1 = f1_score(y_true, y_pred, average='weighted')

    print(f"Precision: {precision:.4f}")
    print(f"Recall: {recall:.4f}")
    print(f"F1 Score: {f1:.4f}")

    return precision, recall, f1


def check_data_distribution(y_train, y_test):
    # 检查数据分布
    print("数据分布检查:")
    print(f"训练集标签分布: {np.bincount(y_train.numpy())}")
    print(f"测试集标签分布: {np.bincount(y_test.numpy())}")
    print(f"标签含义: 0=做空, 1=观望, 2=做多")